package com.spzx.report.utils.embedding;

import com.spzx.report.constant.PineconeConst;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.pinecone.PineconeEmbeddingStore;
import lombok.extern.slf4j.Slf4j;

import java.util.HashMap;
import java.util.List;
import java.util.Map;


/**
 * Package: com.spzx.report.utils
 *
 * @author xiongkunYang
 * @description:
 * @since 2025/6/10 21:15
 */
@Slf4j
public class PineconeSimilarlyUtil {

    public static Map<String, String> getPineconeSimilarityEmbeddings(String question, String index) {

        // 创建向量存储
        EmbeddingStore<TextSegment> embeddingStore = PineconeEmbeddingStore.builder()
                .apiKey(PineconeConst.API_KEY)
                .index(index)
                .metadataTextKey("column")
                .metadataTextKey("word")
                .nameSpace(PineconeConst.NAME_SPACE)
                .build();

        // 先将要拿来匹配的文本向量化
        Embedding queryEmbedding = EmbeddingUtil.getEmbedding(question);

        // 创建搜索请求对象
        EmbeddingSearchRequest searchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                // 匹配最相似的一条记录
                .maxResults(1)
                // 最小匹配阈值
                .minScore(0.8)
                .build();

        // 根据搜索请求 searchRequest 在向量存储中进行相似度搜索
        EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(searchRequest);

        System.out.println(searchResult);
        // searchResult.matches()：获取搜索结果中的匹配项列表。
        // .get(0)：从匹配项列表中获取第一个匹配项
        List<EmbeddingMatch<TextSegment>> matches = searchResult.matches();
        if (!matches.isEmpty()) {
            EmbeddingMatch<TextSegment> embeddingMatch = searchResult.matches().get(0);
            log.info("匹配文本：{}, 相似度分数{}", embeddingMatch.embedded(), embeddingMatch.score());
            // 获取匹配项
            String column = embeddingMatch.embedded().metadata().getString("column");
            String word = embeddingMatch.embedded().metadata().getString("word");

            // 返回结果
            Map<String, String> resultMap = new HashMap<>();
            resultMap.put("column", column);
            resultMap.put("word", word);
            return resultMap;
        } else {
            // 返回结果
            log.error("未找到匹配项, 返回默认数据");
            Map<String, String> resultMap = new HashMap<>();
            resultMap.put("column", "");
            resultMap.put("word", "");
            return resultMap;
        }
    }

}